package com.ruoyi.peach.controller;

import com.alibaba.dashscope.aigc.generation.*;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.util.Collections;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

@RestController
@CrossOrigin
@RequestMapping("/aiDialogue")
public class AiStreamController {

    @Value("${aliyun.api-key}")
    private String apiKey; // 从配置文件读取API Key

    /**
     * 本接口采用两个线程：
     *  1. 一个线程用于与阿里云大模型建立流式连接，实时将产生的小段内容追加到共享字符串中
     *  2. 另一个线程周期性检查共享字符串中是否有新内容（通过 sum 和 index 的比较），将新内容通过 Flux 流式返回给前端
     *
     *  为保证两个线程间数据安全，使用了 synchronized 块对共享字符串的操作进行同步，并用 wait/notify 协调线程唤醒。
     *
     *  变量说明：
     *      - sum：累计从阿里云大模型获取到的内容总长度
     *      - index：后端已传输给前端的内容长度
     *      - arr：存储累计的内容
     *      - sig：标记阿里云数据传输是否结束
     */
    @GetMapping(value = "/send", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> streamGeneration(@RequestParam String prompt) {

        // 使用 Reactor Sinks 创建一个单向的发布者用于 Flux 消息的发送
        Sinks.Many<String> sink = Sinks.many().unicast().onBackpressureBuffer();

        // 共享变量
        AtomicInteger sum = new AtomicInteger(0);       // 当前累计产生的内容长度
        AtomicInteger index = new AtomicInteger(0);     // 当前已传输给前端的内容长度
        StringBuilder arr = new StringBuilder();        // 用于保存产生的内容
        AtomicBoolean sig = new AtomicBoolean(false);   // 用于标记内容是否传输完毕

        // 线程1：与阿里云大模型通信，并将每个返回的小段追加到 arr 中
        new Thread(() -> {
            Generation gen = new Generation();
            Message userMsg = Message.builder()
                    .role(Role.USER.getValue())
                    .content(prompt)
                    .build();

            GenerationParam param = GenerationParam.builder()
                    .apiKey(apiKey)
                    .model("qwen-plus")
                    .messages(Collections.singletonList(userMsg))
                    .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                    .incrementalOutput(true)
                    .build();

            try {
                // 当接收到一个小段内容时，将其追加到 arr 并更新 sum 的值，然后唤醒等待线程
                gen.streamCall(param).blockingForEach(result -> {
                    String content = result.getOutput().getChoices().get(0).getMessage().getContent();

                    synchronized (arr) {
                        arr.append(content);
                        sum.addAndGet(content.length());
                        // 新内容产生，唤醒等待线程
                        arr.notifyAll();
                    }
                });
            } catch (Exception e) {
                sink.tryEmitError(new RuntimeException("AI服务异常：" + e.getMessage()));
            }
            // 阿里云数据传输完毕，置 sig 标志为 true 并唤醒等待线程，通知传输线程处理完毕逻辑
            sig.set(true);
            synchronized (arr) {
                arr.notifyAll();
            }
        }, "AliyunWorker").start();

        // 线程2：定时检查 arr 是否有新内容，将新内容发送给前端
        new Thread(() -> {
            while (!sig.get() || index.get() < sum.get()) {
                synchronized (arr) {
                    // 如果没有新内容但阿里云还未结束传输，则等待一段时间
                    while (index.get() >= sum.get() && !sig.get()) {
                        try {
                            arr.wait(200);
                        } catch (InterruptedException ie) {
                            Thread.currentThread().interrupt();
                        }
                    }
                    // 如果存在新内容，则发送给前端
                    if (index.get() < sum.get()) {
                        String contentToSend = arr.substring(index.get(), sum.get());
                        index.set(sum.get()); // 更新已发送内容的指针
                        String data = contentToSend;
                        sink.tryEmitNext(data);
                        System.out.println("传输给前端内容: " + contentToSend);
                    }
                }
            }
            // 最后完成流式传输
            sink.tryEmitComplete();
        }, "FrontendWorker").start();

        return sink.asFlux();
    }
}
